import torch
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from bert_config_zc import Bert_Config_ZC

# 创建配置对象
conf = Bert_Config_ZC()

def load_raw_data(data_path):
    tmp_df = pd.read_csv(data_path, header=None, names=['label', 'text'])
    res_list = [] #存储的数据格式为[(text, label), (text, label)...]
    for index, row in tmp_df.iterrows():
        label = row['label']
        label = int(label)
        text = row['text']
        text = text.strip() # 安全期间取出前后空格
        res_list.append((text, label))
    return res_list

class PediatricTextDataset(Dataset):
    def __init__(self, data_list):
        super().__init__()
        self.data_list = data_list
        self.data_len = len(self.data_list)

    def __len__(self):
        return self.data_len

    def __getitem__(self, item):
        # 索引异常值处理
        item = min(max(item, 0), len(self.data_list))
        text, label = self.data_list[item]
        return text, label

def collate_fn(batch_data):
    texts_list, labels_list = zip(*batch_data)
    texts_token = conf.bert_tokenizer(texts_list, padding='max_length', max_length=conf.seq_max_len, truncation=True, return_tensors='pt')
    input_ids = texts_token['input_ids']
    attention_mask = texts_token['attention_mask']
    labels = torch.tensor(labels_list)
    return input_ids, attention_mask, labels

def create_dataloader():
    train_data_list = load_raw_data(conf.train_data_path)
    dev_data_list = load_raw_data(conf.dev_data_path)
    test_data_list = load_raw_data(conf.test_data_path)

    train_dataset = PediatricTextDataset(train_data_list)
    dev_dataset = PediatricTextDataset(dev_data_list)
    test_dataset = PediatricTextDataset(test_data_list)

    train_dataloader = DataLoader(train_dataset, batch_size=conf.batch_size, shuffle=True, collate_fn=collate_fn)
    dev_dataloader = DataLoader(dev_dataset, batch_size=conf.batch_size, shuffle=True, collate_fn=collate_fn)
    test_dataloader = DataLoader(test_dataset, batch_size=conf.batch_size, shuffle=True, collate_fn=collate_fn)

    return train_dataloader, dev_dataloader, test_dataloader


if __name__ == '__main__':
    # dev_data_list = load_raw_data(conf.dev_data_path)
    # print(dev_data_list[:10])

    train_dataloader, dev_dataloader, test_dataloader = create_dataloader()
    for idx, (input_ids, attention_mask, labels) in enumerate(dev_dataloader):
        print(f'input_ids：{input_ids}')
        print(f'attention_mask：{attention_mask}')
        print(f'labels：{labels}')
        break



